import nest_asyncio
from dotenv import load_dotenv
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain.chat_models import init_chat_model
from langchain_core.prompts import ChatPromptTemplate

from src.Basic.tools.generatePDF import generate_pdf
from src.Basic.tools.summaryWebSize import summary_website


def main():
    load_dotenv(override=True)

    #定义模型
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    # 定义 天气查询 工具函数
    tools = [summary_website,generate_pdf]
    agent_prompt = ChatPromptTemplate.from_messages(
        [
            ("system", "你是数据处理专家，善于帮助用户处理数据方面的工作"),
            ("user", """
                            请抓取{url}网页的内容并进行总结，优化总结后内容，使其更适合PDF报告格式：
                            请重新组织内容，包括：
                            1.清晰的标题和结构
                            2.要点总结
                            3.详细说明
                            4.使用林妹妹等
                            优化后的内容：
                        """),
            ("placeholder", "{agent_scratchpad}"),
        ])
    agent = create_tool_calling_agent(llm=model, tools=tools, prompt=agent_prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

    while True:
        url = input("请输入要查询的网址:")
        if url == "exit":
            break


        response = agent_executor.invoke({"url": url})
        print(response["output"])

if __name__ == "__main__":
    main()